#!/usr/bin/env python3
"""
Extract NCEP local BUFR definitions (sequence 3-61-223 and dependencies)
from an AMSU-A radiance BUFR file and emit ecCodes-compatible definition files.
"""
import argparse
from pathlib import Path
from typing import Dict, List

from pybufrkit.decoder import Decoder
from pybufrkit.commands import generate_bufr_message
from pybufrkit.descriptors import SequenceDescriptor, FixedReplicationDescriptor

SEQUENCE_IDS = {361223, 361234}
EXTRA_TABLE_B = {12206}  # descriptors missing from stock ecCodes tables
ABBREVIATIONS = {
    12206: "coldSpaceTempCorrection",
}
NAMES = {
    12206: "COLD SPACE TEMPERATURE CORRECTION",
}
CREX_UNIT_DEFAULT = "C"
CREX_SCALE_DEFAULT = 2
CREX_WIDTH_DEFAULT = 4


def descriptor_id(desc) -> int:
    return desc.id


def collect_sequence(node, seq_defs: Dict[int, List[int]], element_meta: Dict[int, object]):
    desc = node.descriptor
    seq_id = descriptor_id(desc)
    if seq_id in seq_defs:
        return

    entries: List[int] = []
    for child in node.members:
        child_desc = child.descriptor
        fid = descriptor_id(child_desc)
        F = child_desc.F

        if F == 3:  # nested sequence
            collect_sequence(child, seq_defs, element_meta)
            entries.append(fid)
        elif F == 1 and isinstance(child_desc, FixedReplicationDescriptor):
            entries.append(fid)
            if child.members:
                first = child.members[0]
                first_desc = first.descriptor
                collect_sequence(first, seq_defs, element_meta)
                entries.append(descriptor_id(first_desc))
        else:
            entries.append(fid)
            if F == 0:
                element_meta.setdefault(fid, child_desc)

    seq_defs[seq_id] = entries


def format_sequence_entry(seq_id: int, entries: List[int]) -> str:
    seq_key = f"{seq_id:06d}"
    values = [f"{val:06d}" for val in entries]
    chunks = []
    line = []
    for idx, val in enumerate(values, 1):
        line.append(val)
        if idx % 6 == 0 and idx != len(values):
            chunks.append(", ".join(line))
            line = []
    if line:
        chunks.append(", ".join(line))
    if len(chunks) == 1:
        body = chunks[0]
    else:
        body = ",\n              ".join(chunks)
    return f'"{seq_key}" = [  {body} ]\n'


def format_element_entry(desc) -> str:
    fid = descriptor_id(desc)
    code = f"{fid:06d}"
    abbreviation = ABBREVIATIONS.get(fid, desc.name.strip().lower().replace(" ", "_"))
    name = NAMES.get(fid, desc.name.strip().replace("  ", " "))
    unit = desc.unit.strip() if isinstance(desc.unit, str) else ""
    scale = desc.scale
    reference = desc.refval
    width = desc.nbits
    return (
        f"{code}|{abbreviation}|double|{name}|{unit}|{scale}|{reference}|{width}|"
        f"{CREX_UNIT_DEFAULT}|{CREX_SCALE_DEFAULT}|{CREX_WIDTH_DEFAULT}\n"
    )


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("bufr_file", type=Path, help="Input AMSU-A BUFR file")
    parser.add_argument(
        "output_dir",
        type=Path,
        help="Directory for ecCodes local table files",
    )
    args = parser.parse_args()

    decoder = Decoder(fallback_or_ignore_missing_tables=True)
    with args.bufr_file.open("rb") as fh:
        data = fh.read()

    sequence_nodes = {}
    element_meta: Dict[int, object] = {}

    for message in generate_bufr_message(
        decoder, data, file_path=str(args.bufr_file), wire_template_data=True
    ):
        if message.n_subsets.value == 0:
            continue
        template = message.sections[3].template_data.value
        template.wire()
        for node in template.decoded_nodes_all_subsets[0]:
            desc = node.descriptor
            if isinstance(desc, SequenceDescriptor) and desc.id in SEQUENCE_IDS:
                collect_sequence(node, sequence_nodes, element_meta)
        if sequence_nodes:
            break

    if not sequence_nodes:
        raise SystemExit("Target sequence 3-61-223 not found in BUFR file")

    args.output_dir.mkdir(parents=True, exist_ok=True)
    seq_path = args.output_dir / "sequence.def"
    with seq_path.open("w", encoding="ascii") as f:
        for seq_id in sorted(sequence_nodes):
            f.write(format_sequence_entry(seq_id, sequence_nodes[seq_id]))

    element_path = args.output_dir / "element.table"
    with element_path.open("w", encoding="ascii") as f:
        f.write("#code|abbreviation|type|name|unit|scale|reference|width|crex_unit|crex_scale|crex_width\n")
        for fid in sorted(EXTRA_TABLE_B):
            desc = element_meta.get(fid)
            if desc is None:
                raise SystemExit(f"Descriptor {fid} not encountered; cannot emit element.table entry")
            f.write(format_element_entry(desc))

    print(f"Wrote {seq_path}")
    print(f"Wrote {element_path}")


if __name__ == "__main__":
    main()
